{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Classification using TAO Classification PyT\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is FAN ?\n",
    "\n",
    "FAN (Fully Attentional Network) is a transformer-based family of backbone from NVIDIA research that achieves SOTA in robustness against various corruptions. This family of backbone can easily generalize to new domains, be more robust to noise, blur etc. Key design behind FAN block is the attentional channel processing module that leads to robust representation learning. FAN can be used for image classification tasks as well as downstream tasks such as object detection and segmentation.\n",
    "FAN can be useful when domain gap exists between the training and testing datasets, for example, a computer vision model is trained using high-resolution images taken in well-lit studio conditions. The training dataset consists of professional photographs with ideal lighting and controlled environments. However, during testing, the model encounters low-resolution images captured by surveillance cameras in outdoor settings with varying lighting conditions and weather effects, so to bridge this domain gap, employing techniques such as FAN can help enhance the model's adaptability to the testing dataset's distinct visual characteristics and challenges.\n",
    "\n",
    "## What is GCViT ?\n",
    "\n",
    "The model in this instance is an image classification model based on [GCViT](https://arxiv.org/abs/2206.09959) architecture. Global context vision transformer (GC ViT),enhances parameter and compute utilization for computer vision. It leverages global context self-attention modules, joint with standard local self-attention, to effectively and efficiently model both long and short-range spatial interactions, without the need for expensive operations such as computing attention masks or shifting local windows."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample prediction of Classification PyT model\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/classification_pyt/sample.jpg?raw=true\" width=\"960\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Train a fan_small_12_p4_hybrid model on the Cats and Dogs dataset\n",
    "* Evaluate the trained model.\n",
    "* Run Inference on the trained model.\n",
    "* Export the trained model to a .onnx file for deployment to DeepStream.\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `classification` model\n",
    "which you may deploy via [Triton](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps)\n",
    "or [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of Classification using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Inferences](#head-6)\n",
    "7. [Deploy](#head-7)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=FIXME\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"classification_pyt\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=/path/to/local/tao-experiments/classification\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "# Point to the 'deps' folder in samples from where you are launching notebook inside classification folder.\n",
    "os.environ[\"PROJECT_DIR\"]=FIXME\n",
    "# Set your encryption key, and use the same key for all commands\n",
    "%env NUM_GPUS = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tao_configs = {\n",
    "   \"Mounts\":[\n",
    "       # Mapping the data directory\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "           \"destination\": \"/workspace/tao-experiments\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       },\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         }\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tao_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Prepare dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the `Cats and Dogs` dataset for the classification tutorial. Please use the following installation steps. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://www.dropbox.com/s/wml49yrtdo53mie/cats_dogs_dataset_reorg.zip?dl=0 -O cats_dogs_dataset.zip\n",
    "!unzip -qo cats_dogs_dataset.zip -d $HOST_DATA_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the following dependencies for running the dataset preparation scripts\n",
    "!pip3 install Cython==0.29.36\n",
    "!pip3 install -r $PROJECT_DIR/deps/requirements-pip.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Verify downloaded dataset <a class=\"anchor\" id=\"head-1-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $HOST_DATA_DIR/cats_dogs_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!ls -l $HOST_DATA_DIR\n",
    "!if [ ! -f $HOST_DATA_DIR/cats_dogs_dataset/classes.txt ]; then echo 'Dataset Not Found, Please Download.'; else echo 'Successfully Found Cats Dogs Dataset.';fi"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-2\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* checkpoint_config: configure the checkpoint setting\n",
    "    * interval: number of iterations at which checkpoint needs to be saved\n",
    "* train_config: configure the training hyperparameters\n",
    "    * optim_config\n",
    "    * epochs\n",
    "    * checkpoint_interval\n",
    "* dataset_config: configure the dataset and augmentation methods\n",
    "    * train_img_dirs\n",
    "    * train_ann_dirs\n",
    "    * pallete: color and mapping class for each class\n",
    "    * output_shape\n",
    "    * batch_size\n",
    "    * workers: number of workers to do data loading\n",
    "    * clips_per_video: number of clips to be sampled from single video\n",
    "    * augmentation_config\n",
    "\n",
    "Please refer to the TAO documentation about Classification to get all the parameters that are configurable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** If you are using the Logistic Regression head, the following parameters from the spec file model config should be used:\n",
    "\n",
    "* model:\n",
    "  * backbone:\n",
    "    * freeze: true\n",
    "    * pretrained: \"/path/to/NV_DINOV2_518.pth\"\n",
    "  * head:\n",
    "    * lr_head:\n",
    "      * C: 0.316   # tunable\n",
    "      * max_iter: 5000   # tunable\n",
    "    * type: LogisticRegressionHead\n",
    "    * num_classes: 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/train_cats_dogs.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Download pre-trained model <a class=\"anchor\" id=\"head-1-4\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use NGC CLI to get the pre-trained models. For more details, go to ngc.nvidia.com and click the SETUP on the navigation bar. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "import os\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $HOST_RESULTS_DIR/ngccli\n",
    "\n",
    "# # Remove any previously existing CLI installations\n",
    "!rm -rf $HOST_RESULTS_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $HOST_RESULTS_DIR/ngccli\n",
    "!unzip -u \"$HOST_RESULTS_DIR/ngccli/$CLI\" -d $HOST_RESULTS_DIR/ngccli/\n",
    "!rm $HOST_RESULTS_DIR/ngccli/*.zip\n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"HOST_RESULTS_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/pretrained_fan_classification_imagenet:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_PROJECT_DIR/pretrained_fan_hybrid_small/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/pretrained_fan_classification_imagenet:fan_hybrid_small --dest $LOCAL_PROJECT_DIR/pretrained_fan_hybrid_small"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $LOCAL_PROJECT_DIR/pretrained_fan_hybrid_small/pretrained_fan_classification_imagenet_vfan_hybrid_small"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# This is the suitable number of epochs for this model with pretrained weights. Please change this value as needed.\n",
    "%env EPOCHS = 3\n",
    "\n",
    "print(\"Train Classification Model\")\n",
    "!tao model classification_pyt train \\\n",
    "                  -e $SPECS_DIR/train_cats_dogs.yaml \\\n",
    "                  results_dir=$RESULTS_DIR/classification_experiment \\\n",
    "                  train.num_gpus=$NUM_GPUS \\\n",
    "                  model.init_cfg.checkpoint=/workspace/tao-experiments/pretrained_fan_hybrid_small/pretrained_fan_classification_imagenet_vfan_hybrid_small/fan_hybrid_small.pth \\\n",
    "                  train.train_config.runner.max_epochs=$EPOCHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"To resume from a checkpoint, use the below command. Update the epoch number accordingly\")\n",
    "# !tao model classification_pyt train \\\n",
    "#                   -e $SPECS_DIR/train_cats_dogs.yaml \\\n",
    "#                   results_dir$RESULTS_DIR/classification_experiment \\\n",
    "#                   train.num_gpus=$NUM_GPUS \\\n",
    "#                   train.train_config.resume_training_checkpoint_path=$RESULTS_DIR/classification_experiment/train/epoch_20.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('PyTorch checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/classification_experiment/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "%env NUM_EPOCH=3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Rename a model: Note that the training is not deterministic, so you may change the model name accordingly.')\n",
    "print('---------------------')\n",
    "# NOTE: The following command may require `sudo`. You can run the command outside the notebook.\n",
    "!find $HOST_RESULTS_DIR/classification_experiment/train -name \"epoch_$NUM_EPOCH.pth\" | xargs realpath | xargs -I {} mv {} $HOST_RESULTS_DIR/classification_experiment/train/classification_model.pth \n",
    "!ls -ltrh $HOST_RESULTS_DIR/classification_experiment/train/classification_model.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-4\"></a>\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluate Cats Dogs Classification Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model classification_pyt evaluate \\\n",
    "                    -e $SPECS_DIR/test_cats_dogs.yaml \\\n",
    "                    evaluate.checkpoint=$RESULTS_DIR/classification_experiment/train/classification_model.pth \\\n",
    "                    results_dir=$RESULTS_DIR/classification_experiment"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Inferences <a class=\"anchor\" id=\"head-5\"></a>\n",
    "In this section, we run the classification inference tool to generate inferences with the trained classification models and print the results. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model classification_pyt inference \\\n",
    "                    -e $SPECS_DIR/test_cats_dogs.yaml \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/classification_experiment/train/classification_model.pth \\\n",
    "                    results_dir=$RESULTS_DIR/classification_experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the results\n",
    "!cat $HOST_RESULTS_DIR/classification_experiment/inference/result.csv"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the inference with images from the csv file. It contains the following columns - Image Name, class_label, class_confidence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install Deps\n",
    "!pip3 install pillow\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image \n",
    "import os\n",
    "import csv\n",
    "from math import ceil\n",
    "import random \n",
    "\n",
    "DATA_DIR = os.environ.get('HOST_DATA_DIR')\n",
    "DATA_DOWNLOAD_DIR = os.environ.get('DATA_DIR')\n",
    "RESULT_DIR = os.environ.get('HOST_RESULTS_DIR')\n",
    "csv_path = os.path.join(RESULT_DIR, \"classification_experiment/inference/\" 'result.csv')\n",
    "results = []\n",
    "with open(csv_path) as csv_file:\n",
    "    csv_reader = csv.reader(csv_file, delimiter=',')\n",
    "    for row in csv_reader:\n",
    "        results.append((row[0], row[1]))\n",
    "random.shuffle(results)\n",
    "\n",
    "w,h = 200,200\n",
    "fig = plt.figure(figsize=(30,30))\n",
    "columns = 5\n",
    "rows = 1\n",
    "for i in range(1, columns*rows + 1):\n",
    "    ax = fig.add_subplot(rows, columns,i)\n",
    "    img = Image.open(results[i][0].replace(DATA_DOWNLOAD_DIR, DATA_DIR))\n",
    "    img = img.resize((w,h), Image.ANTIALIAS)\n",
    "    plt.imshow(img)\n",
    "    ax.set_title(results[i][1], fontsize=40)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Deploy! <a class=\"anchor\" id=\"head-6\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export the Classification model to ONNX model\n",
    "# NOTE: Export is done on single GPU - GPU num need not be provided\n",
    "\n",
    "!tao model classification_pyt export \\\n",
    "                   -e $SPECS_DIR/export_cats_dogs.yaml \\\n",
    "                   export.checkpoint=$RESULTS_DIR/classification_experiment/train/classification_model.pth \\\n",
    "                   export.onnx_file=$RESULTS_DIR/classification_experiment/export/classification_model_export.onnx \\\n",
    "                   results_dir=$RESULTS_DIR/classification_experiment/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a TensorRT Engine using TAO Deploy\n",
    "!tao deploy classification_pyt gen_trt_engine \\\n",
    "                   -e $SPECS_DIR/export_cats_dogs.yaml \\\n",
    "                   gen_trt_engine.onnx_file=$RESULTS_DIR/classification_experiment/export/classification_model_export.onnx \\\n",
    "                   gen_trt_engine.trt_engine=$RESULTS_DIR/classification_experiment/gen_trt_engine/classification_model_export.engine \\\n",
    "                   results_dir=$RESULTS_DIR/classification_experiment/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run evaluation using the generated TensorRT Engine\n",
    "!tao deploy classification_pyt evaluate \\\n",
    "                   -e $SPECS_DIR/test_cats_dogs.yaml \\\n",
    "                   evaluate.trt_engine=$RESULTS_DIR/classification_experiment/gen_trt_engine/classification_model_export.engine \\\n",
    "                   results_dir=$RESULTS_DIR/classification_experiment/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run inference using the generated TensorRT Engine\n",
    "!tao deploy classification_pyt inference \\\n",
    "                   -e $SPECS_DIR/test_cats_dogs.yaml \\\n",
    "                   inference.trt_engine=$RESULTS_DIR/classification_experiment/gen_trt_engine/classification_model_export.engine \\\n",
    "                   results_dir=$RESULTS_DIR/classification_experiment/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the results\n",
    "!cat $HOST_RESULTS_DIR/classification_experiment/trt_inference/result.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize Inference\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image \n",
    "import os\n",
    "import csv\n",
    "from math import ceil\n",
    "import random \n",
    "\n",
    "DATA_DIR = os.environ.get('HOST_DATA_DIR')\n",
    "DATA_DOWNLOAD_DIR = os.environ.get('DATA_DIR')\n",
    "RESULT_DIR = os.environ.get('HOST_RESULTS_DIR')\n",
    "csv_path = os.path.join(RESULT_DIR, \"classification_experiment/trt_inference/\" 'result.csv')\n",
    "results = []\n",
    "with open(csv_path) as csv_file:\n",
    "    csv_reader = csv.reader(csv_file, delimiter=',')\n",
    "    for row in csv_reader:\n",
    "        results.append((row[0], row[1]))\n",
    "random.shuffle(results)\n",
    "\n",
    "w,h = 200,200\n",
    "fig = plt.figure(figsize=(30,30))\n",
    "columns = 5\n",
    "rows = 1\n",
    "for i in range(1, columns*rows + 1):\n",
    "    ax = fig.add_subplot(rows, columns,i)\n",
    "    img = Image.open(results[i][0].replace(DATA_DOWNLOAD_DIR, DATA_DIR))\n",
    "    img = img.resize((w,h), Image.ANTIALIAS)\n",
    "    plt.imshow(img)\n",
    "    ax.set_title(results[i][1], fontsize=40)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  },
  "vscode": {
   "interpreter": {
    "hash": "8dd506d74b0e800f607ce4c2aea2a7d91bf4a69ac59ceef67481b8a20350139b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
